import os
from data_loader import Dataset
import matplotlib.pyplot as plt
from typing import List
from metrics import metric
import numpy as np
import pandas as pd

root_list = ['Brent', 'RCLC', 'WTI', 'NYMEX', 'OPEC', '中国油价']
data_list = ['Brenta', 'RCLCa', 'WTIa', 'NYMEX天然气期货每年', 'OPEC原油产量', '中国油价数据每年']

OPTION = {
    'SINGLE': 0,
    'TV': 1,
    'TVT': 2
}


def savePlotResults(option: int, normalize=True, save_pdf=True,
                    save_csv=False, is_predict=True, model=None, **kwargs):
    root_path = os.path.join('./dataset', root_list[option])
    data_path = data_list[option] + '.csv'
    if os.path.exists(f'./results/{root_list[option]}') is False:
        os.makedirs(f'./results/{root_list[option]}')
    dataset = Dataset(root_path, data_path)

    data = dataset.data
    data_raw = dataset.inverse_transform(data)

    if is_predict is False:
        if normalize:
            plot_data(data, figsize=(20, 10), label='price', xlabel='Year', xticks=dataset.year)
            if save_pdf:
                plt.savefig(f'./results/{root_list[option]}/{data_list[option]}(标准化).pdf')
        else:
            plot_data(data_raw, figsize=(20, 10), label='price', xlabel='Year', xticks=dataset.year)
            if save_pdf:
                plt.savefig(f'./results/{root_list[option]}/{data_list[option]}.pdf')
        plt.show()
    else:
        _test_model(model)

        time_steps = kwargs['pred_steps']
        train = model.dataset.data.flatten()
        index = np.arange(len(train) + time_steps)
        xtick = np.arange(model.dataset.year[0], model.dataset.year[-1] + 1 + time_steps)

        train_inverse = model.dataset.inverse_transform(train.reshape(-1, 1)).flatten()
        preds = model.predict(time_steps)
        preds_inverse = model.dataset.inverse_transform(preds.reshape(-1, 1)).flatten()

        _plot_predict(train, np.concatenate([train, preds]), index, xtick, option, time_steps, normalized=True)
        _plot_predict(train_inverse, np.concatenate([train_inverse, preds_inverse]), index, xtick, option, time_steps,
                      normalized=False)

        if save_csv:
            result = np.around(np.concatenate([train_inverse, preds_inverse]), 2)
            dtframe = np.concatenate([xtick.astype('str').reshape(-1, 1), result.reshape(-1, 1)], axis=-1)
            dtframe = pd.DataFrame(dtframe, columns=['date', f'price({kwargs["units"]})'])
            pd.DataFrame.to_csv(dtframe, f"./results/{root_list[option]}/predict_{time_steps}.csv")
            return dtframe


def _test_model(model):
    train = np.concatenate([model.dataset.train_data.flatten(), model.dataset.val_data.flatten()])
    test = model.dataset.test_data.flatten()
    time_steps = len(test)
    preds = model.predict(time_steps=time_steps)
    plt.plot(np.concatenate([train, preds]), label='preds')
    plt.plot(np.concatenate([train, test]), label='true')
    plt.legend()
    plt.show()
    mae, mse, rmse, mape, mspe = metric(preds, test)
    verbose = list(zip(('mae', 'mse', 'rmse', 'mape', 'mspe'), (mae, mse, rmse, mape, mspe)))
    print(verbose)


def _plot_predict(train, predict, index, xtick, option,
                  time_steps, rotation=30, normalized=False):
    plt.figure(figsize=(15, 5))
    plt.plot(predict, label='predict')
    plt.plot(train, label='known')
    plt.xticks(index, xtick, rotation=rotation)
    plt.legend()
    if normalized:
        plt.savefig(f"./results/{root_list[option]}/predict_{time_steps}(normalized).pdf")
    else:
        plt.savefig(f"./results/{root_list[option]}/predict_{time_steps}.pdf")
    plt.show()


def plot_data(data: np.ndarray, figsize: tuple, label: str, xlabel: str, xticks):
    plt.figure(figsize=figsize)
    index = np.arange(len(data))
    plt.plot(index, data, label=label)
    plt.xlabel(xlabel, loc='right')
    plt.xticks(index, xticks, rotation=-30)
    plt.legend()


def split_plot(aim: str, figsize: tuple, data, label=None):
    if OPTION[aim] == 0:
        plot_data(data, figsize, label)
    elif OPTION[aim] == 1:
        plt.figure(figsize=figsize)
        l1 = len(data[0])
        l2 = len(data[1])
        plt.plot(np.arange(l1), data[0], label='train')
        plt.plot(np.arange(l1, l1 + l2), data[1], label='val')
        plt.legend()
    elif OPTION[aim] == 2:
        datas1 = np.concatenate(data, axis=0)
        datas2 = np.concatenate(data[1:], axis=0)
        plt.figure(figsize=figsize)
        l1 = len(data[0])
        l2 = len(data[1])
        l3 = len(data[2])
        plt.plot(datas1, c='b')
        plt.plot(np.arange(l1), data[0], label='train', c='b')
        plt.plot(np.arange(l1, l1 + l2 + l3), datas2, label='train', c='orange')
        plt.plot(np.arange(l1, l1 + l2), data[1], label='val', c='orange')
        plt.plot(np.arange(l1 + l2, l1 + l2 + l3), data[2], label='test', c='r')
        plt.legend()
    else:
        raise NotImplementedError


def ConcateFrame(frame_list: List, pred_step=10):
    min_year = [l.year.min() for l in frame_list]
    max_year = [l.year.max() for l in frame_list]
    min_year = max(min_year)
    max_year = max(max_year)
    frame_list[-1].inv_data = np.concatenate([frame_list[-1].inv_data, np.zeros((pred_step, 1))], axis=0)
    range = max_year - min_year + 1
    index = np.arange(min_year, max_year + 1).astype(str)
    data = np.concatenate([l.inv_data[-range:] for l in frame_list], axis=-1).reshape(-1, 6)
    data = np.concatenate([index.reshape(-1, 1), data], axis=-1)
    column = ['date'] + root_list
    frame = pd.DataFrame(data=data, columns=column)
    pd.DataFrame.to_csv(frame, './dataset/all.csv', index=False)
    return frame


if __name__ == '__main__':
    # savePlotResults(3, normalize=True, save_pdf=True, save_csv=False, is_predict=False)
    frames = [Dataset(root_path=os.path.join('./results', root_list[option]),
                      data_path='predict_10.csv')
              for option in range(5)]
    frames.append(Dataset(root_path=os.path.join('./dataset', root_list[-1]),
                          data_path='中国油价数据每年.csv'))
    frame = ConcateFrame(frames)
